package com.classactionpl.jaas.activemq;

import java.util.HashSet;
import java.util.Set;
import java.util.concurrent.CopyOnWriteArrayList;

import org.apache.activemq.broker.Broker;
import org.apache.activemq.broker.ConnectionContext;
import org.apache.activemq.broker.Connector;
import org.apache.activemq.broker.jmx.ManagedTransportConnector;
import org.apache.activemq.command.ConnectionInfo;
import org.apache.activemq.jaas.GroupPrincipal;
import org.apache.activemq.security.JaasCertificateAuthenticationBroker;
import org.apache.activemq.security.SecurityContext;
import org.apache.activemq.transport.tcp.SslTransportServer;
import org.apache.commons.logging.Log;
import org.apache.commons.logging.LogFactory;

/**
 * SSLCertificateAuthenticationBroker is a variation of
 * JaasCertificateAuthenticationBroker where only the SSL transport causes an
 * authentication against an X.509 certificate.
 * 
 * Where a non-SSL transport is used a username and set of group principals
 * passed into the constructor are used. This then permits any authorisation
 * policies to be satisfied.
 * 
 * @author huntc
 * 
 */
public class SslCertificateAuthenticationBroker extends
		JaasCertificateAuthenticationBroker {

	private Log logger = LogFactory
			.getLog(SslCertificateAuthenticationBroker.class);

	private String nonsslUsername;
	private Set<GroupPrincipal> nonsslGroups;

	private final CopyOnWriteArrayList<SecurityContext> securityContexts = new CopyOnWriteArrayList<SecurityContext>();

	/**
	 * @param next
	 * @param jaasConfiguration
	 * @param nonsslUsername
	 *            The username to use internally. Must be distinct from other
	 *            usernames in the system.
	 * @param nonsslGroupNames
	 *            A set of string describing which groups the nonssl user should
	 *            receive. These groups can then be used subsequently for
	 *            authorisation checks.
	 */
	public SslCertificateAuthenticationBroker(Broker next,
			String jaasConfiguration, String nonsslUsername,
			Set<String> nonsslGroupNames) {
		super(next, jaasConfiguration);

		if (logger.isDebugEnabled()) {
			logger.debug("nonsslUsername: " + nonsslUsername
					+ " - nonsslGroups: " + nonsslGroupNames);
		}

		this.nonsslUsername = nonsslUsername;

		this.nonsslGroups = new HashSet<GroupPrincipal>(nonsslGroupNames.size());

		for (String nonsslGroupName : nonsslGroupNames) {
			this.nonsslGroups.add(new GroupPrincipal(nonsslGroupName));
		}
	}

	/*
	 * (non-Javadoc)
	 * 
	 * @seeorg.apache.activemq.security.JaasCertificateAuthenticationBroker#
	 * addConnection(org.apache.activemq.broker.ConnectionContext,
	 * org.apache.activemq.command.ConnectionInfo)
	 */
	@Override
	public void addConnection(ConnectionContext context, ConnectionInfo info)
			throws Exception {

		if (context.getSecurityContext() == null) {
			boolean isSSL;
			Connector connector = context.getConnector();
			if (connector instanceof ManagedTransportConnector) {
				ManagedTransportConnector managedTransportConnector = (ManagedTransportConnector) connector;
				isSSL = (managedTransportConnector.getServer() instanceof SslTransportServer);
			} else {
				isSSL = false;
			}

			if (!isSSL) {
				if (logger.isDebugEnabled()) {
					logger.debug("nonssl transport detected - using user: "
							+ nonsslUsername);
				}

				info.setUserName(nonsslUsername);
				info.setPassword("");

				SecurityContext securityContext = new SecurityContext(
						nonsslUsername) {
					public Set<?> getPrincipals() {
						return nonsslGroups;
					}
				};

				context.setSecurityContext(securityContext);

				securityContexts.add(securityContext);

			} else {
				if (logger.isDebugEnabled()) {
					logger.debug("ssl transport detected - using certs");
				}
			}
		}

		super.addConnection(context, info);
	}

	/*
	 * (non-Javadoc)
	 * 
	 * @seeorg.apache.activemq.security.JaasCertificateAuthenticationBroker#
	 * removeConnection(org.apache.activemq.broker.ConnectionContext,
	 * org.apache.activemq.command.ConnectionInfo, java.lang.Throwable)
	 */
	@Override
	public void removeConnection(ConnectionContext context,
			ConnectionInfo info, Throwable error) throws Exception {
		super.removeConnection(context, info, error);
		if (securityContexts.remove(context.getSecurityContext())) {
			context.setSecurityContext(null);
		}
	}

}
